/****************************************************************
 * @file   : cmat.c
 * @author : NingJian (freegodly@gmail.com)
 * @date   : 06/07/21 16:22:14
 * @brief  :
 ****************************************************************/

#include "cmat.h"
#include <math.h> 
#include <stdio.h>
#include <string.h>

///////////////////////////////  private fun start
static void cmatprint_dim(CMat *ma, int dim, size_t start_offset);
static void cmatcopy_x(CMat *src, int *src_info, CMat *des, int *des_info, int start_dim);
static void cmatcopy_y(CMat *src, int *src_info, CMat *des, int *des_info, int start_dim, int repeat_count, int old_dim_size);
static void cmatmult_x(CMat *mat_a, int *mat_a_dim_size, CMat *mat_b, int *mat_b_dim_size, int start_dim_a, int start_dim_b);
///////////////////////////////  private fun end

CMat *CMat_create(int dim, int *dim_size)
{
  CMat *ma = calloc(1, sizeof(CMat));

  ma->dim = dim;

  ma->data_len = 1;
  ma->dim_size = calloc(dim, sizeof(int));
  int i = 0;
  for (i = 0; i < dim; ++i)
  {
    ma->dim_size[i] = dim_size[i];
    ma->data_len *= dim_size[i];
  }

  ma->data_ptr = calloc(ma->data_len, sizeof(double));

  return ma;
}

CMat *CMat_clone(CMat *mat)
{
  CMat *new_mat = CMat_create(mat->dim, mat->dim_size);
  memcpy(new_mat->data_ptr, mat->data_ptr, mat->data_len * sizeof(double));
  return new_mat;
}

void CMat_delete(CMat *ma)
{
  if (!ma)
    return;
  free(ma->dim_size);
  free(ma->data_ptr);
  free(ma);
}

void cmatprint_dim(CMat *ma, int dim, size_t start_offset)
{
  if (dim < ma->dim - 1)
  {
    size_t dim_size = 1;
    int j = 0;
    for (j = dim + 1; j < ma->dim; ++j)
    {
      dim_size *= ma->dim_size[j];
    }
    int i = 0;

    for (j = 0; j < ma->dim_size[dim]; ++j)
    {
      for (i = 0; i < dim + 1; ++i)
        printf("  ");
      printf("\r\n");
      for (i = 0; i < dim + 1; ++i)
        printf("  ");
      printf("[");
      cmatprint_dim(ma, dim + 1, start_offset + dim_size * j);
      printf("],\r\n");
    }
    for (i = 0; i < dim; ++i)
      printf("  ");
  }
  else
  {
    size_t dim_size = ma->dim_size[dim];
    size_t i = 0;
    for (i = 0; i < dim_size; ++i)
    {
      printf("%f", ma->data_ptr[start_offset + i]);
      if (i != dim_size - 1)
        printf(",");
    }
  }
}

void CMat_view(CMat *mat, int new_dim, int *new_dim_info)
{
  int *new_dim_size = calloc(new_dim + 1, sizeof(int));
  int i = 0;
  for (i = 0; i < new_dim; ++i)
  {
    new_dim_size[i] = new_dim_info[i];
  }
  free(mat->dim_size);
  mat->dim_size = new_dim_size;
  mat->dim = new_dim;
}

void CMat_print(CMat *ma)
{
  if (!ma)
    return;

  printf("dim = %d,data_len = %d, dim_size = [", ma->dim, (int)ma->data_len);
  int i = 0;
  for (i = 0; i < ma->dim; ++i)
  {
    printf("%d", ma->dim_size[i]);
    if (i != ma->dim - 1)
      printf(",");
    else
      printf("]\r\n");
  }

  printf("[\r\n");
  cmatprint_dim(ma, 0, 0);
  printf("]\r\n");
}

double CMat_get(CMat *ma, int *dim_info)
{
  size_t offset = CMat_diminfo2index(ma, dim_info);

  return ma->data_ptr[offset];
}

void CMat_set(CMat *ma, int *dim_info, double value)
{
  size_t offset = CMat_diminfo2index(ma, dim_info);
  ma->data_ptr[offset] = value;
}

void CMat_index2diminfo(CMat *ma, size_t index, int *dim_info)
{
  int i = 0;
  int j = 0;
  size_t index_tmp = index;
  for (i = 0; i < ma->dim - 1; ++i)
  {
    int tmp = 1;
    for (j = i + 1; j < ma->dim; ++j)
    {
      tmp *= ma->dim_size[j];
    }
    dim_info[i] = index_tmp / tmp;
    index_tmp = index_tmp - dim_info[i] * tmp;
  }
  dim_info[ma->dim - 1] = index_tmp;
}

size_t CMat_diminfo2index(CMat *ma, int *dim_info)
{
  size_t offset = 0;

  int j = 0;
  int i = 0;
  for (i = 0; i < ma->dim; ++i)
  {
    int tmp = dim_info[i];
    for (j = i + 1; j < ma->dim; ++j)
    {
      tmp *= ma->dim_size[j];
    }
    offset += tmp;
  }

  return offset;
}

CMat *CMat_stack(CMat *x, CMat *y, int dim)
{
  // now only support dim = -1
  // x y for 2 dim
  if (dim != -1 || x->dim != y->dim)
    return NULL;
  int *new_dim_size = calloc(x->dim + 1, sizeof(int));
  size_t r = 0;
  size_t c = 0;
  int i = 0;
  for (i = 0; i < x->dim; ++i)
  {
    new_dim_size[i] = x->dim_size[i];
  }
  new_dim_size[x->dim] = 2;
  CMat *out = CMat_create(x->dim + 1, new_dim_size);

  int *old_dim_size = calloc(x->dim, sizeof(int));

  for (r = 0; r < x->dim_size[0]; ++r)
  {
    old_dim_size[0] = r;
    new_dim_size[0] = r;
    for (c = 0; c < x->dim_size[1]; ++c)
    {
      old_dim_size[1] = c;
      new_dim_size[1] = c;

      new_dim_size[2] = 0;
      CMat_set(out, new_dim_size, CMat_get(x, old_dim_size));
      new_dim_size[2] = 1;
      CMat_set(out, new_dim_size, CMat_get(y, old_dim_size));
    }
  }

  free(new_dim_size);
  free(old_dim_size);

  return out;
}

CMat *CMat_arange_unsqueeze_repeat(int size, int direction)
{
  int dim_info[2];
  dim_info[0] = size;
  dim_info[1] = size;
  CMat *ma = CMat_create(2, dim_info);

  int r = 0;
  int c = 0;
  for (r = 0; r < size; ++r)
  {
    dim_info[0] = r;
    for (c = 0; c < size; ++c)
    {
      dim_info[1] = c;
      if (direction == 0)
      { // 行方向涨
        CMat_set(ma, dim_info, c);
      }
      else
      {
        CMat_set(ma, dim_info, r);
      }
    }
  }

  return ma;
}

void CMat_unsqueeze(CMat *ma, int dim)
{
  int dim_old = ma->dim;
  int *dim_size_old = ma->dim_size;

  int *new_dim_size = calloc(dim_old + 1, sizeof(int));
  int i = 0;
  int j = 0;
  for (i = 0; i < dim_old; ++i, ++j)
  {
    if (j != dim)
    {
      new_dim_size[j] = dim_size_old[i];
    }
    else
    {
      new_dim_size[j] = 1;
      --i;
    }
  }
  free(dim_size_old);
  ma->dim_size = new_dim_size;
  ma->dim += 1;
}

void cmatcopy_x(CMat *src, int *src_info, CMat *des, int *des_info, int start_dim)
{
  int i = 0;
  if (start_dim < src->dim - 1)
  {
    for (i = 0; i < src->dim_size[start_dim]; ++i)
    {
      src_info[start_dim] = i;
      des_info[start_dim] = i;
      cmatcopy_x(src, src_info, des, des_info, start_dim + 1);
    }
  }
  else if (start_dim == src->dim - 1)
  {
    for (i = 0; i < src->dim_size[start_dim]; ++i)
    {
      src_info[start_dim] = i;
      des_info[start_dim] = i;

      CMat_set(des, des_info, CMat_get(src, src_info));
    }
  }
  else
  {
    CMat_set(des, des_info, CMat_get(src, src_info));
  }
}

void cmatmult_x(CMat *mat_a, int *mat_a_dim_size, CMat *mat_b, int *mat_b_dim_size, int start_dim_a, int start_dim_b)
{
  if (start_dim_a < mat_a->dim - 1)
  {
    int i = 0;
    for (i = 0; i < mat_a->dim_size[start_dim_a]; ++i)
    {
      mat_a_dim_size[start_dim_a] = i;
      mat_b_dim_size[start_dim_b] = i;
      cmatmult_x(mat_a, mat_a_dim_size, mat_b, mat_b_dim_size, start_dim_a + 1, start_dim_b + 1);
    }
  }
  else if (start_dim_a == mat_a->dim - 1)
  {
    int i = 0;
    for (i = 0; i < mat_a->dim_size[start_dim_a]; ++i)
    {
      mat_a_dim_size[start_dim_a] = i;
      mat_b_dim_size[start_dim_b] = i;
      //  printf("start_dim_a %d , value = %d  start_dim_b = %d   value = %d \r\n", start_dim_a, mat_a->dim_size[start_dim_a],start_dim_b, mat_b_dim_size[start_dim_b]);
      //CMat_set(mat_a, mat_a_dim_size,1);
      //  printf("mat_a_dim_size [%d,%d,%d,%d,%d]  mat_b_dim_size [%d,%d]\r\n",
      //        mat_a_dim_size[0],
      //        mat_a_dim_size[1],
      //        mat_a_dim_size[2],
      //        mat_a_dim_size[3],
      //        mat_a_dim_size[4],
      //        mat_b_dim_size[0],
      //        mat_b_dim_size[1]);
      CMat_set(mat_a, mat_a_dim_size, CMat_get(mat_a, mat_a_dim_size) * CMat_get(mat_b, mat_b_dim_size));
    }
  }
  else
  {
    CMat_set(mat_a, mat_a_dim_size, CMat_get(mat_a, mat_a_dim_size) * CMat_get(mat_b, mat_b_dim_size));
  }
}

void cmatcopy_y(CMat *src, int *src_info, CMat *des, int *des_info, int start_dim, int repeat_count, int old_dim_size)
{
  // for repeat
  size_t j = 0;
  size_t k = 0;
  size_t n = 0;
  size_t m = 0;
  size_t l = 0;
  size_t i = 0;

  for (i = 0; i < src->dim; ++i)
  {
    src_info[i] = 0;
    des_info[i] = 0;
  }

  i = start_dim;

  for (j = 1; j < repeat_count; ++j) //需要copy
  {
    size_t left_change_count = 1;
    for (l = 0; l < i; ++l)
    {
      left_change_count *= src->dim_size[l];
    }
    //printf("left_change_count : %d \r\n", left_change_count);

    for (m = 0; m < left_change_count; ++m)
    {
      if (i > 0) // index i = 0的时候，left没有数据
      {
        int index_tmp = m;
        for (l = 0; l < i - 1; ++l)
        {
          int tmp = 1;
          for (n = l + 1; n < i; ++n)
          {
            tmp *= src->dim_size[n];
          }
          src_info[l] = index_tmp / tmp;
          des_info[l] = index_tmp / tmp;
          index_tmp = index_tmp - src_info[l] * tmp;
        }
        src_info[i - 1] = index_tmp;
        des_info[i - 1] = index_tmp;
      }

      for (k = 0; k < old_dim_size; ++k)
      {
        src_info[i] = k;
        des_info[i] = k + j * old_dim_size;

        // printf("old [%d,%d,%d,%d,%d] ,  new [%d,%d,%d,%d,%d]\r\n",
        //        src_info[0],
        //        src_info[1],
        //        src_info[2],
        //        src_info[3],
        //        src_info[4],
        //        des_info[0],
        //        des_info[1],
        //        des_info[2],
        //        des_info[3],
        //        des_info[4]);
        cmatcopy_x(src, src_info, des, des_info, i + 1);
        for (l = i; l < src->dim; ++l)
        {
          src_info[l] = 0;
          des_info[l] = 0;
        }
      }
    }
  }
}

CMat *CMat_repeat(CMat *ma, int *repeat_info)
{
  // repeat_info length must eq ma dim length

  // calc new size
  int *new_dim_size = calloc(ma->dim, sizeof(int));
  int *old_dim_size = calloc(ma->dim, sizeof(int));
  size_t i = 0;
  for (i = 0; i < ma->dim; ++i)
  {
    new_dim_size[i] = repeat_info[i] * ma->dim_size[i];
  }

  CMat *out = CMat_create(ma->dim, new_dim_size);

  //copy old
  for (i = 0; i < ma->data_len; ++i)
  {
    CMat_index2diminfo(ma, i, new_dim_size);
    CMat_set(out, new_dim_size, CMat_get(ma, new_dim_size));
  }

  // copy ext
  for (i = 0; i < ma->dim; ++i)
  {
    if (repeat_info[i] > 1) // need copy
    {
      cmatcopy_y(out, old_dim_size, out, new_dim_size, i, repeat_info[i], ma->dim_size[i]);
    }
  }

  free(new_dim_size);
  free(old_dim_size);

  return out;
}

void CMat_sigmoid(CMat *mat)
{
  int i = 0;
  for (i = 0; i < mat->data_len; ++i)
  {
    mat->data_ptr[i] = 1.0 / (1 + exp(-1.0 * mat->data_ptr[i]));
  }
}

void CMat_exp(CMat *mat)
{
  int i = 0;
  for (i = 0; i < mat->data_len; ++i)
  {
    mat->data_ptr[i] = exp(mat->data_ptr[i]);
  }
}

void CMat_mult_value(CMat *mat, double value)
{
  int i = 0;
  for (i = 0; i < mat->data_len; ++i)
  {
    mat->data_ptr[i] *= value;
  }
}

void CMat_mult_mat(CMat *mat_a, CMat *mat_b)
{
  // require mat_a dim >=   mat_b dim
  int i = 0;
  if (mat_a->dim != mat_b->dim)
  {
    int *mat_a_dim_size = calloc(mat_a->dim, sizeof(int));
    int *mat_b_dim_size = calloc(mat_b->dim, sizeof(int));

    size_t left_dim = mat_a->dim - mat_b->dim;
    size_t left_change_count = 1;
    size_t m = 0;
    size_t l = 0;
    size_t n = 0;

    for (i = 0; i < left_dim; ++i)
    {
      left_change_count *= mat_a->dim_size[i];
    }
    //printf("left_change_count : %d \r\n", left_change_count);
    for (m = 0; m < left_change_count; ++m)
    {
      if (i > 0) // index i = 0的时候，left没有数据
      {
        int index_tmp = m;
        for (l = 0; l < left_dim - 1; ++l)
        {
          int tmp = 1;
          for (n = l + 1; n < left_dim; ++n)
          {
            tmp *= mat_a->dim_size[n];
          }
          mat_a_dim_size[l] = index_tmp / tmp;
          index_tmp = index_tmp - mat_a_dim_size[l] * tmp;
        }
        mat_a_dim_size[i - 1] = index_tmp;
      }
      // printf("mat_a_dim_size [%d,%d,%d,%d,%d]  mat_b_dim_size [%d,%d]\r\n",
      //        mat_a_dim_size[0],
      //        mat_a_dim_size[1],
      //        mat_a_dim_size[2],
      //        mat_a_dim_size[3],
      //        mat_a_dim_size[4],
      //        mat_b_dim_size[0],
      //        mat_b_dim_size[1]);
      cmatmult_x(mat_a, mat_a_dim_size, mat_b, mat_b_dim_size, left_dim, 0);
      for (l = left_dim; l < mat_a->dim; ++l)
      {
        mat_a_dim_size[l] = 0;
      }
      for (l = 0; l < mat_b->dim; ++l)
      {
        mat_b_dim_size[l] = 0;
      }
    }
    free(mat_a_dim_size);
    free(mat_b_dim_size);
  }
  else
  {
    for (i = 0; i < mat_a->data_len; ++i)
    {
      mat_a->data_ptr[i] *= mat_b->data_ptr[i];
    }
  }
}

CMat *CMat_permute(CMat *mat, int *dim_info)
{
  int *new_dim_size = calloc(mat->dim, sizeof(int));
  int i = 0;
  for (i = 0; i < mat->dim; ++i)
  {
    new_dim_size[i] = mat->dim_size[dim_info[i]];
  }

  CMat *new_mat = CMat_create(mat->dim, new_dim_size);

  // copy data
  int *mat_dim_size = calloc(mat->dim, sizeof(int));

  size_t j = 0;
  for (j = 0; j < mat->data_len; ++j)
  {
    CMat_index2diminfo(mat, j, mat_dim_size);
    // convert to new dim_size
    for (i = 0; i < mat->dim; ++i)
    {
      new_dim_size[i] = mat_dim_size[dim_info[i]];
    }
    CMat_set(new_mat, new_dim_size, mat->data_ptr[j]);
  }

  free(new_dim_size);
  free(mat_dim_size);

  return new_mat;
}

CMat *CMat_pick(CMat *mat, int (*pick_dim_info)[2])
{
  int *new_dim_size = calloc(mat->dim, sizeof(int));
  int i = 0;
  for (i = 0; i < mat->dim; ++i)
  {
    new_dim_size[i] = pick_dim_info[i][1] - pick_dim_info[i][0];
  }
  CMat *new_mat = CMat_create(mat->dim, new_dim_size);

  // copy data
  int *mat_dim_size = calloc(mat->dim, sizeof(int));
  size_t j = 0;
  for (j = 0; j < new_mat->data_len; ++j)
  {
    CMat_index2diminfo(new_mat, j, new_dim_size);
    // convert to old dim_size
    for (i = 0; i < mat->dim; ++i)
    {
      mat_dim_size[i] = new_dim_size[i] + pick_dim_info[i][0];
    }
    CMat_set(new_mat, new_dim_size, CMat_get(mat, mat_dim_size));
  }

  free(new_dim_size);
  free(mat_dim_size);
  return new_mat;
}

void CMat_add_mat(CMat *mat_a, CMat *mat_b)
{
  int i = 0;
  for (i = 0; i < mat_a->data_len; ++i)
  {
    mat_a->data_ptr[i] += mat_b->data_ptr[i];
  }
}

CMat *CMat_cat(CMat **mats, int mats_len, int dim)
{
  // only support dim = -1
  int *new_dim_size = calloc(mats[0]->dim, sizeof(int));

  int *new_dim_len = calloc(mats_len, sizeof(int));

  int i = 0;
  int n = 0;

  for (i = 0; i < mats[0]->dim; ++i)
  {
    new_dim_size[i] += mats[0]->dim_size[i];
  }

  new_dim_len[0] = mats[0]->dim_size[mats[0]->dim - 1];
  for (n = 1; n < mats_len; ++n)
  {
    new_dim_size[mats[0]->dim - 1] += mats[n]->dim_size[mats[0]->dim - 1];
    new_dim_len[n] = new_dim_len[n - 1] + mats[n]->dim_size[mats[0]->dim - 1];
  }

  CMat *new_mat = CMat_create(mats[0]->dim, new_dim_size);

  int *mat_dim_size = calloc(mats[0]->dim, sizeof(int));

  size_t j = 0;
  for (j = 0; j < new_mat->data_len; ++j)
  {
    CMat_index2diminfo(new_mat, j, new_dim_size);
    // convert to old dim_size
    for (i = 0; i < mats[0]->dim; ++i)
    {
      mat_dim_size[i] = new_dim_size[i];
    }

    int last_dim_index = new_dim_size[mats[0]->dim - 1];
    int copy_from_matn = 0;
    for (n = 0; n < mats_len; ++n)
    {
      if (new_dim_len[n] > last_dim_index)
      {
        copy_from_matn = n;
        break;
      }
    }

    if (copy_from_matn > 0)
    {
      mat_dim_size[mats[0]->dim - 1] -= new_dim_len[copy_from_matn -1];
    }

    // printf("new [%d,%d,%d,%d,%d]  [matn = %d] old [%d,%d,%d,%d,%d]\r\n",
    //        new_dim_size[0],
    //        new_dim_size[1],
    //        new_dim_size[2],
    //        new_dim_size[3],
    //        new_dim_size[4],
    //        copy_from_matn,
    //        mat_dim_size[0],
    //        mat_dim_size[1],
    //        mat_dim_size[2],
    //        mat_dim_size[3],
    //        mat_dim_size[4]);

    CMat_set(new_mat, new_dim_size, CMat_get(mats[copy_from_matn], mat_dim_size));
  }
  free(new_dim_size);
  free(new_dim_len);
  free(mat_dim_size);
  return new_mat;
}